#pragma once
#include <cuda_runtime.h>

#include <cmath>
#include <cassert>
#include <cstdint>
#include <cstring>  // for memset
#include <limits>
#include <queue>
#include <utility>  // for std::pair

#include "common/data.h"

class ESDFMap {
public:
    ESDFMap(int map_row_size, int map_col_size, float resolution);
    ~ESDFMap();

    void reset();
    void init_occupancy_boundary(const std::vector<Polygon>& polygon_vec);
    void init_occupancy_inside(const std::vector<Polygon>& polygon_vec);
    void calc_esdf_cpu();
    void calc_esdf_gpu();

    inline uint8_t get_occupancy(int row, int col) const {
        assert(is_index_valid(row, col));
        return occupancy_buffer_inflate_[to_address(row, col)];
    }

    inline float get_distance_cpu(int row, int col) const {
        assert(is_index_valid(row, col));
        return distance_buffer_[to_address(row, col)];
    }

    inline float get_distance_cpu_with_sign(int row, int col) const {
        assert(is_index_valid(row, col));
        float dist = get_distance_cpu(row, col);
        return get_occupancy(row, col) == static_cast<uint8_t>(OccupancyType::INSIDE) ? -dist : dist;
    }

    inline float get_distance_gpu(int row, int col) const {
        assert(is_index_valid(row, col));
        return distance_buffer_gpu_res_[to_address(row, col)];
    }

    inline float get_distance_gpu_with_sign(int row, int col) const {
        assert(is_index_valid(row, col));
        float dist = get_distance_gpu(row, col);
        return get_occupancy(row, col) == static_cast<uint8_t>(OccupancyType::INSIDE) ? -dist : dist;
    }

    bool is_index_valid(int row, int col) const { return row < map_row_size_ && row >= 0 && col < map_col_size_ && col >= 0; }
    int to_address(int row, int col) const { return row * map_col_size_ + col; }
    int float_to_grid(float value) { return static_cast<int>(std::round(value / resolution_)); }

    void debug_print_occupancy() const;
    void debug_print_esdf_cpu() const;
    void debug_print_esdf_gpu() const;

private:
    // 在计算机图形学中，将连续的直线离散化为像素点的过程称为直线光栅化。
    // 最常用的算法是Bresenham 算法，它通过整数运算高效地确定直线经过的像素点。
    void bresenham_line(float x0, float y0, float x1, float y1);
    void bfs_set_occupancy(const Point2d& internal_point, OccupancyType flag);

    // ref to https://github.com/HKUST-Aerial-Robotics/Fast-Planner/blob/master/fast_planner/plan_env/src/sdf_map.cpp
    template <typename F_get_val, typename F_set_val>
    void fill_esdf(F_get_val f_get_val, F_set_val f_set_val, int start, int end, int dim_size) {
        int v[dim_size];
        float z[dim_size + 1];

        int k = start;
        v[start] = start;
        z[start] = -std::numeric_limits<float>::max();
        z[start + 1] = std::numeric_limits<float>::max();

        for (int q = start + 1; q <= end; q++) {
            k++;
            float s;

            do {
                k--;
                s = ((f_get_val(q) + q * q) - (f_get_val(v[k]) + v[k] * v[k])) / (2 * q - 2 * v[k]);
            } while (s <= z[k]);

            k++;

            v[k] = q;
            z[k] = s;
            z[k + 1] = std::numeric_limits<float>::max();
        }

        k = start;

        for (int q = start; q <= end; q++) {
            while (z[k + 1] < q) k++;
            float val = (q - v[k]) * (q - v[k]) + f_get_val(v[k]);
            f_set_val(q, val);
        }
    }

private:
    int map_row_size_;
    int map_col_size_;
    float resolution_;

    uint8_t* occupancy_buffer_inflate_ = nullptr;
    float* tmp_buffer_ = nullptr;
    float* distance_buffer_ = nullptr;

    int* obs_index_data_ = nullptr;        // 第i个元素内容v, 表示第i个障碍物像素点所在列的位置
    int* obs_start_index_data_ = nullptr;  // 第j个元素内容index, 表示j行中第一个障碍物在obs_index_data_中所在的位置

    cudaStream_t stream_;
    uint8_t* occupancy_buffer_inflate_gpu_ = nullptr;
    float* tmp_buffer_gpu_ = nullptr;
    float* distance_buffer_gpu_ = nullptr;

    int* obs_index_data_gpu_ = nullptr;
    int* obs_start_index_data_gpu_ = nullptr;

    float* distance_buffer_gpu_res_ = nullptr;
};